I'm currently skipping the argmin argmax tests because of #3463 but that should change.

So we can't take the product of an array with negative numbers? The logarithm will just NaN the output

awkward/src/awkward/_connect/jax/reducers.py

Lines 296 to 299 in c93da2f

	# See issue https://github.com/google/jax/issues/9296
	result = jax.numpy.exp(
	jax.ops.segment_sum(jax.numpy.log(array.data), parents.data)
	)

So we can't take the product of an array with negative numbers? The logarithm will just NaN the output

I may have fixed that in 9874810. I still haven't made ak.any to behave like the cpu case when segments have zeros though.

@ikrommyd - impressive! There are only 5 tests that fails: 3 for ak.any, 1 for ak.count, and 1 for ak.sum. It looks like the latter is failing with a boolean dtype. Perhaps, return should make an int?

So I made ci pass but there are a few things that are to be done for sure

1. The code needs refactoring, I don't like how it looks at all, it's very hacky
2. The reducer tests should definitely test more input array cases
3. We need to make sure I'm not breaking something that should work but is untested due to the smaller amount of tests for the jax backend